package ink.mint.activation;

import ink.mint.matrix.Matrix;
import ink.mint.matrix.Matrixable;

import java.io.Serializable;

public class ActivationFunction implements Activateable, Serializable {
    /**
     * 激活类型
     */
    private ActivationType activationType;
    /**
     * 矩阵
     */
    private Matrixable matrix;

    /**
     * 无参构造激活函数
     */
    public ActivationFunction() {
        this.matrix = new Matrix();
    }

    /**
     * 激活函数的构造
     *
     * @param activationType 激活类型
     */
    public ActivationFunction(ActivationType activationType) {
        this.matrix = new Matrix();
        this.activationType = activationType;
    }

    /**
     * 激活函数解
     *
     * @param z              输入z
     * @param activationType 激活类型
     * @return 激活函数的解
     */
    public double ActivationSolve(double z, ActivationType activationType) {
        switch (activationType) {
            case Sigmoid:
                return sigmoid(z);
            case RELU:
                return relu(z);
            case Softmax:
                return softmax(z);
            case Tanh:
                return tanh(z);
            default:
                return 0;
        }
    }

    /**
     * 激活函数解
     *
     * @param z              输入z
     * @param activationType 激活类型
     * @return 激活函数的解
     */
    public double[] ActivationSolve(double[] z, ActivationType activationType) {
        int zr = z.length;
        switch (activationType) {
            case Sigmoid:
                return sigmoid(z);
            case RELU:
                return relu(z);
            case Softmax:
                return softmax(z);
            case Tanh:
                return tanh(z);
            default:
                return new double[zr];
        }
    }

    /**
     * 激活函数解
     *
     * @param z              输入z
     * @param activationType 激活类型
     * @return 激活函数的解
     */
    public double[][] ActivationSolve(double[][] z, ActivationType activationType) {
        int zr = z.length;
        int zc = z[0].length;
        switch (activationType) {
            case Sigmoid:
                return sigmoid(z);
            case RELU:
                return relu(z);
            case Softmax:
                return softmax(z);
            case Tanh:
                return tanh(z);
            default:
                return new double[zr][zc];
        }
    }

    /**
     * 求导后的激活函数解
     *
     * @param z              输入z
     * @param activationType 激活类型
     * @return 激活函数的求导解
     */
    public double derivationActivationSolve(double z, ActivationType activationType) {
        switch (activationType) {
            case Sigmoid:
                return derivationSigmoid(z);
            case RELU:
                return derivationRelu(z);
            case Softmax:
                return derivationSoftmax(z);
            case Tanh:
                return derivationTanh(z);
            default:
                return 0;
        }
    }

    /**
     * 求导后的激活函数解
     *
     * @param z              输入z
     * @param activationType 激活类型
     * @return 激活函数的求导解
     */
    public double[] derivationActivationSolve(double[] z, ActivationType activationType) {
        int zr = z.length;
        switch (activationType) {
            case Sigmoid:
                return derivationSigmoid(z);
            case RELU:
                return derivationRelu(z);
            case Softmax:
                return derivationSoftmax(z);
            case Tanh:
                return derivationTanh(z);
            default:
                return new double[zr];
        }
    }

    /**
     * 求导后的激活函数解
     *
     * @param z              输入z
     * @param activationType 激活类型
     * @return 激活函数的求导解
     */
    public double[][] derivationActivationSolve(double[][] z, ActivationType activationType) {
        int zr = z.length;
        int zc = z[0].length;
        switch (activationType) {
            case Sigmoid:
                return derivationSigmoid(z);
            case RELU:
                return derivationRelu(z);
            case Softmax:
                return derivationSoftmax(z);
            case Tanh:
                return derivationTanh(z);
            default:
                return new double[zr][zc];
        }
    }

    /**
     * sigmoid 函数
     *
     * @param z 函数输入值
     * @return sigmoid函数解
     */
    @Override
    public double sigmoid(double z) {
        return 1 / (1 + Math.exp(-1 * z));
    }

    /**
     * sigmoid 函数
     *
     * @param z 输入函数的向量
     * @return sigmoid函数向量解
     */
    @Override
    public double[] sigmoid(double[] z) {
        int zr = z.length;
        double[] result = new double[zr];
        for (int i = 0; i < zr; i++) {
            result[i] = 1 / (1 + Math.exp(-1 * z[i]));
        }
        return result;
    }

    /**
     * sigmoid 函数
     *
     * @param z 输入函数的矩阵
     * @return sigmoid函数矩阵解
     */
    @Override
    public double[][] sigmoid(double[][] z) {
        int zr = z.length;
        int zc = z[0].length;
        double[][] result = new double[zr][zc];
        for (int i = 0; i < zr; i++) {
            for (int j = 0; j < zc; j++) {
                result[i][j] = 1 / (1 + Math.exp(-1 * z[i][j]));
            }
        }
        return result;
    }

    /**
     * 求导后的sigmoid函数
     *
     * @param z 函数输入
     * @return 函数的解
     */
    @Override
    public double derivationSigmoid(double z) {
        double a = sigmoid(z);
        return a * (1 - a);
    }

    /**
     * 求导后的sigmoid函数
     *
     * @param z 输入函数的向量
     * @return 函数的向量解
     */
    @Override
    public double[] derivationSigmoid(double[] z) {
        int zr = z.length;
        double[] a = sigmoid(z);
        double[] result = new double[zr];
        for (int i = 0; i < zr; i++) {
            result[i] = a[i] * (1 - a[i]);
        }
        return result;
    }

    /**
     * 求导后的sigmoid函数
     *
     * @param z 输入函数的矩阵
     * @return 函数的矩阵解
     */
    @Override
    public double[][] derivationSigmoid(double[][] z) {
        int zr = z.length;
        int zc = z[0].length;
        double[][] a = sigmoid(z);
        double[][] result = new double[zr][zc];
        for (int i = 0; i < zr; i++) {
            for (int j = 0; j < zc; j++) {
                result[i][j] = a[i][j] * (1 - a[i][j]);
            }
        }
        return result;
    }

    /**
     * relu函数
     *
     * @param z 函数输入的值
     * @return relu函数解
     */
    @Override
    public double relu(double z) {
        return Math.max(0, z);
    }

    /**
     * relu函数
     *
     * @param z 输入函数的向量值
     * @return relu函数向量解
     */
    @Override
    public double[] relu(double[] z) {
        int zr = z.length;
        double[] result = new double[zr];
        for (int i = 0; i < zr; i++) {
            result[i] = Math.max(0, z[i]);
        }
        return result;
    }

    /**
     * relu函数
     *
     * @param z 输入函数的矩阵值
     * @return relu函数矩阵解
     */
    @Override
    public double[][] relu(double[][] z) {
        int zr = z.length;
        int zc = z[0].length;
        double[][] result = new double[zr][zc];
        for (int i = 0; i < zr; i++) {
            for (int j = 0; j < zc; j++) {
                result[i][j] = Math.max(0, z[i][j]);
            }
        }
        return result;
    }

    /**
     * 求导后的relu函数
     *
     * @param z 输入函数的值
     * @return 求导后的relu函数解
     */
    @Override
    public double derivationRelu(double z) {
        if (z >= 0) {
            return 1;
        } else {
            return 0;
        }
    }

    /**
     * 求导后的relu函数
     *
     * @param z 输入函数的向量值
     * @return 求导后的relu函数向量解
     */
    @Override
    public double[] derivationRelu(double[] z) {
        int zr = z.length;
        double[] result = new double[zr];
        for (int i = 0; i < zr; i++) {
            if (z[i] >= 0) {
                result[i] = 1;
            } else {
                result[i] = 0;
            }
        }
        return result;
    }

    /**
     * 求导后的relu函数
     *
     * @param z 输入函数的矩阵值
     * @return 求导后的relu函数矩阵解
     */
    @Override
    public double[][] derivationRelu(double[][] z) {
        int zr = z.length;
        int zc = z[0].length;
        double[][] result = new double[zr][zc];
        for (int i = 0; i < zr; i++) {
            for (int j = 0; j < zc; j++) {
                if (z[i][j] >= 0) {
                    result[i][j] = 1;
                } else {
                    result[i][j] = 0;
                }
            }
        }
        return result;
    }

    /**
     * softmax函数
     *
     * @param z 输入函数的值
     * @return softmax函数的解
     */
    @Override
    public double softmax(double z) {
        return 1;
    }

    /**
     * softmax函数
     *
     * @param z 输入函数的向量值
     * @return softmax函数的向量解
     */
    @Override
    public double[] softmax(double[] z) {
        int zr = z.length;
        double[] result = new double[zr];
        double[] softNum = new double[zr];
        for (int i = 0; i < zr; i++) {
            softNum[i] = Math.pow(Math.E, z[i]);
        }
        double softSum = matrix.sum(softNum);
        for (int i = 0; i < zr; i++) {
            result[i] = softNum[i] / softSum;
        }
        return result;
    }

    /**
     * softmax函数
     *
     * @param z 输入函数的矩阵值
     * @return softmax函数的矩阵解
     */
    @Override
    public double[][] softmax(double[][] z) {
        int zr = z.length;
        int zc = z[0].length;
        double[][] result = new double[zr][zc];
        double[][] softNum = new double[zr][zc];
        for (int i = 0; i < zr; i++) {
            double[] row = z[i];
            double[] softmaxRow = softmax(row); // 复用向量版本的实现
            result[i] = softmaxRow;
        }
        return result;
    }

    /**
     * 求导后的softmax函数
     *
     * @param z 输入函数的值
     * @return 求导后的softmax函数解
     */
    @Override
    public double derivationSoftmax(double z) {
        return 0;
    }

    /**
     * 求导后的softmax函数
     *
     * @param z 输入函数的向量值
     * @return 求导后的softmax函数向量解
     */
    @Override
    public double[] derivationSoftmax(double[] z) {
        double[] a = softmax(z);
        int zr = z.length;
        double[] result = new double[zr];
        for (int i = 0; i < zr; i++) {
            result[i] = a[i] * (1 - a[i]);
            for (int j = 0; j < zr; j++) {
                if (i != j) {
                    result[i] -= a[i] * a[j];
                }
            }
        }
        return result;
    }


    /**
     * 求导后的softmax函数
     *
     * @param z 输入函数的矩阵值
     * @return 求导后的softmax函数矩阵解
     */
    @Override
    public double[][] derivationSoftmax(double[][] z) {
        int zr = z.length;
        int zc = z[0].length;
        double[][] result = new double[zr][zc];
        for (int i = 0; i < zr; i++) {
            double[] softPow = matrix.pow(softmax(z[i]), 2);
            result[i] = matrix.sub(softmax(z[i]), softPow);
        }
        return result;
    }

    /**
     * tanh函数
     *
     * @param z 输入函数的值
     * @return tanh函数解
     */
    @Override
    public double tanh(double z) {
        return Math.tanh(z);
    }

    /**
     * tanh函数
     *
     * @param z 输入函数的向量值
     * @return tanh函数向量解
     */
    @Override
    public double[] tanh(double[] z) {
        int zr = z.length;
        double[] result = new double[zr];
        for (int i = 0; i < zr; i++) {
            result[i] = Math.tanh(z[i]);
        }
        return result;
    }

    /**
     * tanh函数
     *
     * @param z 输入函数的矩阵值
     * @return tanh函数矩阵解
     */
    @Override
    public double[][] tanh(double[][] z) {
        int zr = z.length;
        int zc = z[0].length;
        double[][] result = new double[zr][zc];
        for (int i = 0; i < zr; i++) {
            for (int j = 0; j < zc; j++) {
                result[i][j] = Math.tanh(z[i][j]);
            }
        }
        return result;
    }

    /**
     * 求导后的tanh函数
     *
     * @param z 输入函数的值
     * @return 求导后的tanh函数解
     */
    @Override
    public double derivationTanh(double z) {
        return 1 / Math.pow(Math.cosh(z), 2);
    }

    /**
     * 求导后的tanh函数
     *
     * @param z 输入函数的向量值
     * @return 求导后的tanh函数向量解
     */
    @Override
    public double[] derivationTanh(double[] z) {
        int zr = z.length;
        double[] result = new double[zr];
        for (int i = 0; i < zr; i++) {
            result[i] = 1 / Math.pow(Math.cosh(z[i]), 2);
        }
        return result;
    }

    /**
     * 求导后的tanh函数
     *
     * @param z 输入函数的矩阵值
     * @return 求导后的tanh函数矩阵解
     */
    @Override
    public double[][] derivationTanh(double[][] z) {
        int zr = z.length;
        int zc = z[0].length;
        double[][] result = new double[zr][zc];
        for (int i = 0; i < zr; i++) {
            for (int j = 0; j < zc; j++) {
                result[i][j] = 1 / Math.pow(Math.cosh(z[i][j]), 2);
            }
        }
        return result;
    }

}
